'''
Author: zhanwei xu
Date: 2023-05-27 18:50:03
LastEditors: zhanwei xu
LastEditTime: 2024-03-01 20:08:31
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''
import time

from PIL import Image
from hashlib import md5
from fastapi import APIRouter, Request
import requests
import httpx
from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
import json
from .utils import async_request, get_api_keys, get_api_key, check_traffic, generate_image_poe, get_prompt_outurl, generate_image_dall_e,check_traffic_3,request_anytext
from httpx  import Timeout
import os
import io
import random
import aiohttp
from aiohttp import ClientTimeout
import string
from io import BytesIO

def generate_random_string(length):
    characters = string.ascii_letters + string.digits
    random_string = ''.join(random.choice(characters) for _ in range(length))
    return random_string
import asyncio
import functools
import uuid
from concurrent.futures import ThreadPoolExecutor
global img_url
img_url = {}
draw_router = APIRouter()
api_keys = get_api_keys("api_keys.txt")
global api_index 
api_index = random.randint(0, len(api_keys) - 1)
@draw_router.get("/ge_key_value")
async def get_key_value(request: Request):
    # 获取前端请求的数据
    request_data = await request.json()
    try:
        key = request_data["key"]
    except:
        key = "aadadafds"
    key_list = os.listdir("keys")
    if key not in key_list:
        return {'error':'key is not valid'}
    else:
        key_path = "keys/" + key
        with open(key_path, 'r') as f:
            key_conunt = f.read()
            key_conunt = int(key_conunt.strip())
            return {'key_value':key_conunt}

def get_clients(file_path = 'p_b_value.txt'):
    
    f = open(file_path, 'r')
    lines = f.readlines()
    f.close()
    line_index = random.randint(0, len(lines) - 1)
    token = lines[line_index].strip()
    return token
def translate_azure(query):
    endpoint = "https://api.cognitive.microsofttranslator.com/"

    # location, also known as region.
    # required if you're using a multi-service or regional (not global) resource. It can be found in the Azure portal on the Keys and Endpoint page.
    location = "eastus"

    path = '/translate'
    constructed_url = endpoint + path

    params = {
        'api-version': '3.0',
        'from': 'zh-Hans',
        'to': 'en'
    }

    # You can pass more than one object in body.
    body = [{
        'text': query
    }]
    key = get_clients('translate_key.txt')
    headers = {
        'Ocp-Apim-Subscription-Key': key,
        # location required if you're using a multi-service or regional (not global) resource.
        'Ocp-Apim-Subscription-Region': location,
        'Content-type': 'application/json',
        'X-ClientTraceId': str(uuid.uuid4())
    }


    request = requests.post(constructed_url, params=params, headers=headers, json=body)
    count = 0 
    if request.status_code != 200 and count < 5:
        
        key = get_clients('translate_key.txt')
        headers = {
            'Ocp-Apim-Subscription-Key': key,
            # location required if you're using a multi-service or regional (not global) resource.
            'Ocp-Apim-Subscription-Region': location,
            'Content-type': 'application/json',
            'X-ClientTraceId': str(uuid.uuid4())
        }


        request = requests.post(constructed_url, params=params, headers=headers, json=body)
    if request.status_code == 200:
        response = request.json()
        # print(response)
        text = response[0]['translations'][0]['text']
        return text
        
def translate(query):
    
    # Set your own appid/appkey.
    appid = 'yourid'
    appkey = 'yourkey'

    # For list of language codes, please refer to `https://api.fanyi.baidu.com/doc/21`
    from_lang = 'en'
    to_lang =  'zh'

    endpoint = 'http://api.fanyi.baidu.com'
    path = '/api/trans/vip/translate'
    url = endpoint + path

    # query = '一位贵妇人'

    # Generate salt and sign
    def make_md5(s, encoding='utf-8'):
        return md5(s.encode(encoding)).hexdigest()

    salt = random.randint(32768, 65536)
    sign = make_md5(appid + query + str(salt) + appkey)

    # Build request
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    payload = {'appid': appid, 'q': query, 'from': "auto", 'to': "en", 'salt': salt, 'sign': sign}

    # Send request
    r = requests.post(url, params=payload, headers=headers)
    result = r.json()
    dst = result["trans_result"][0]["dst"]

    return dst
def get_prompt(response):
    prompt = ''
    for event in response.iter_content(chunk_size=None):     
            response_str = event.decode('utf-8')  # convert bytes to string
            strs = response_str.split('data: ')[1:]
            for str_ in strs:
                try:
                    response_dict = json.loads(str_)  # convert string to dictionary
                    if "choices" in response_dict and response_dict["choices"]:
                        text = response_dict["choices"][0]["delta"].get("content")
                        if text:
                            text = text.replace('\\n','\n').replace('\\t','\t')
                            
                            text = text.replace('\\n','\n')
                            prompt += text
                except:
                    pass
    return prompt
async def async_post(url, data=None, json = None, headers=None, timeout=10):
    
    loop = asyncio.get_running_loop()
    request = functools.partial(requests.post, url, data=data,json = json)
    try:
        with ThreadPoolExecutor() as executor:
            response = await loop.run_in_executor(
                executor, request
            )
    except:
        # 返回一个空的response，状态码为404
        response = requests.Response()
        response.status_code = 500
    return response
async def async_get(url):
    loop = asyncio.get_running_loop()
    request = functools.partial(requests.get, url)
    try:
        with ThreadPoolExecutor() as executor:
            response = await loop.run_in_executor(
                executor, request
            )
    except:
        # 返回一个空的response，状态码为404
        response = requests.Response()
        response.status_code = 500
    return response

async def get_oss_image(id: str):
    url_zaiwen = "https://user.chatcns.com/asset/get"
    params = {"id": id}
    async with aiohttp.ClientSession() as session:
        async with session.get(url_zaiwen, params=params) as resp:
            data = await resp.json()

            url_image = data['info']['url']
            async with session.get(url_image) as resp2:
                if resp2.status == 200:
                    image_data = await resp2.read()
                    image = Image.open(BytesIO(image_data))
                    return image
                else:
                    print('Failed to retrieve the image')

@draw_router.post("/img_generate2")
async def forward_img_generate2(request: Request):
    global img_url
    '''
    request:{
        "model_name":str 'sdxl_turbo' or 'stabldiffusionxl'
        "prompt":str
        "negative_prompt":str
        "ratio":str '方', '高', '宽'
        "original_image":image file
        "mask_image":image file
        "key": str, optional
        "picture_weight":float, optional
        "height":int, optional
        "width":int, optional
        "step":int,opentional
    }
    '''

    # 获取前端请求的数据
    content_type = request.headers.get("Content-Type")
    if "application/json" in content_type:
        # 处理JSON数据
        print("json")
        request_data = await request.json()
    elif "multipart/form-data" in content_type:
        # 处理表单数据
        print("form")
        request_data = await request.form()
    else:
        return JSONResponse(content={"message": "Unsupported data format"},status_code=400)

    flag = 0
    client_ip = request.client.host
        
    global api_index

    try:
        model_name = request_data["model_name"] # "normal","structure","anime"
    except:
        model_name = "sdxl-turbo"
    try:
        key = request_data["key"]
    except:
        key = "aadadafds"
    
    key_list = os.listdir("keys")
    
    if model_name.startswith("dall"):
        if check_traffic_3(client_ip):
            if (key not in key_list) or (os.stat("keys/" + key).st_size == 0):
                response = {'task_id':"hasexceeded"}
                return JSONResponse(response)
            else:
                flag = 1
    else:
        if check_traffic(client_ip):
            if (key not in key_list) or (os.stat("keys/" + key).st_size == 0):
                response = {'status': {"code":429, 'message': '每24小时只能使用100次，请明天再来,谢谢'}}
                return JSONResponse(response)
            else:
                flag = 1
                
    if flag:
        key_path = "keys/" + key
        # print(key_path)
        # key_path = 'keys/mklrv3rcjpfkzqqh5xgl4iarq8q6ctdp'
        f= open(key_path, 'r')
        key_count = f.read()
        key_count = int(key_count.strip())
        
        model_name = request_data["model_name"] # "normal","structure","anime"
        if model_name.startswith("dall"):
            key_count = key_count-10
        else:
            key_count = key_count -5
        # print(key_count)
        f.close()
        f = open(key_path, 'w')
        # print(key_count)
        try:
            f.write(str(key_count))
        except:
            f.write(str(key_count))
        f.close()
        if key_count <= 0:
            os.remove(key_path)


    prompt = request_data["prompt"]
    try:
        negative_prompt = request_data["negative_prompt"]
    except:
        negative_prompt = ""
    model_name = request_data["model_name"]
    trans_flag = 1
    if model_name.startswith("dall") or model_name.startswith("anytext"):
        trans_flag = 0

    ratio = request_data["ratio"]


    original_image = None
    mask_image = None
    picture_weight = None
    try:
        original_image = request_data['original_image']
        picture_weight = request_data['picture_weight']
        mask_image = request_data['mask_image']
    except:
        pass
    mask_image = request_data.get("mask_image")

    try:
        height = request_data["height"]
        if isinstance(height,str):
            height = None
    except:
        height = None
    try:
        width = request_data["width"]
        if isinstance(width,str):
            width = None
    except:
        width = None
    try:
        step = request_data["step"]
        if isinstance(step,str):
            step = None
    except:
        step = None
    

    
    if isinstance(picture_weight,str):
        picture_weight = float(picture_weight)
    if picture_weight:
        if picture_weight>1:
            picture_weight = picture_weight / 100.0
    
    print(f"height:{height} width:{width} step:{step} picture_weight:{picture_weight} orignal {original_image}")

    if trans_flag:
        loop = asyncio.get_running_loop()
        prompt_ = await loop.run_in_executor(None, translate_azure, prompt)
        if prompt_ is not None:
            prompt = prompt_
    if model_name.startswith( "poe_model"):
        # 截断prompt到只有30个单词
        prompt = prompt.split(' ')
        prompt = prompt[:30]
        prompt = ' '.join(prompt)
        prompt = prompt.replace('\n','')
        client_token = get_clients()
        mode = model_name.split('_')[-1]
        url = await generate_image_poe(client_token,prompt,mode=mode)
        task_id =  'org-'+generate_random_string(10)
        img_url[task_id] = url
        response = {'task_id': task_id}
        return response
    if model_name.startswith("dall"):
        api_key = get_clients('api_keys.txt')
        url = await generate_image_dall_e(prompt,api_key,model=ratio)
        if '安全' in url:
            response = {'status': {"code":429, "message":'Dall-e模型有十分严格的安全检查，您输入的提示词安全检查没有通过'}}
            return JSONResponse(response)

        task_id =  'org-'+generate_random_string(10)
        img_url[task_id] = url
        response = {'task_id': task_id}
#        with open('img_url.json','a') as f:
 #           json_data = {"prompt":prompt,"url":url}
  #          json.dump(json_data,f)
   #         f.write('\n')
        return response


    default_negative_prompt = "NSFW, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, \
        normal quality, jpeg artifacts,signature, watermark, username, blurry, artist name, pixelated, incorrect proportions, bad shading, glitch, wrong colors, incorrect perspective, missing features, over-exaggerated features, distorted, low definition, unwanted objects, misaligned elements, poorly lit, visible grid lines, overly saturated, noticeable compression, unnecessary filters, digital noise"
    loop = asyncio.get_running_loop()
    if negative_prompt:
        negative_prompt = await loop.run_in_executor(None, translate_azure, negative_prompt)
    else:
        negative_prompt = ""
    negative_prompt = negative_prompt + default_negative_prompt

    height_width_json ={
        '方':{
            'height':1024,
            'width':1024,
        },
        '高':{
            'height':1024,
            'width':768,
        },
        '宽':{
            'height':768,
            'width':1024,
        },
        '16:9':{
            'height':720,
            'width':1280,
        },
        '9:16':{
            'height':1280,
            'width':720,
        },
        '无':{
            'height':None,
            'width':None,
        },
    }
    
    if height:
        height = (height // 64) *  64
    if width:
        width = (width // 64) *  64

    if height and not width:
        width = (height_width_json[ratio]['width'] / height_width_json[ratio]['height']) * height
    elif width and not height:
        height = (height_width_json[ratio]['height'] / height_width_json[ratio]['width']) * width
    elif not height and not width:
        height = height_width_json[ratio]['height']
        width = height_width_json[ratio]['width']


    # 不用translate
        
    # if model_name.startswith("anytext"):
    #     data_json = {
    #         "prompt": prompt,
    #         "ddim_steps": step if step else 20,
    #         "image_width": height,
    #         "image_height": width,
    #         "a_prompt": "",
    #         "n_prompt": negative_prompt
    #     }
    #     task_id =  'anytext-'+generate_random_string(10)
    #     asyncio.create_task(request_anytext(task_id, data_json,img_url))
    #     response = {'task_id': task_id}
    #     return response


    # 创建 FormData 对象
    form_data = aiohttp.FormData()
    
    form_data.add_field('prompt', prompt)
    form_data.add_field('negative_prompt', negative_prompt)
    form_data.add_field('model_name', model_name)
    form_data.add_field('height', str(height)) if height else None
    form_data.add_field('width', str(width)) if width else None
    form_data.add_field('step', str(step)) if step else None
    if picture_weight:
        form_data.add_field('picture_weight', str(picture_weight))
    # print('注意注意')
    # print(original_image)
    # 对于文件，使用add_field将文件添加到表单数据中
    print("我没进来")
    if original_image:
        if original_image != 'undefined':
            print("我进来了")
            # 确保光标回到文件的开头
            if isinstance(original_image,str):
                img = await get_oss_image(original_image)
            else:
                original_image.file.seek(0)
                img = Image.open(original_image.file) 

            width, height = img.size
            # 确定最大边长，并计算缩放比例
            max_side_length = 1024
            if width > max_side_length or height > max_side_length:
                # 计算缩放比例
                scale_factor = min(max_side_length / width, max_side_length / height)
                
                # 按比例计算新的大小
                new_width = int(width * scale_factor)
                new_height = int(height * scale_factor)
                
                # 缩放图片
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            # 假设 img 是一个打开的 RGBA 模式的 PIL Image 对象
            if img.mode == 'RGBA':
                # 创建一个白色背景的图片
                background = Image.new('RGB', img.size, (255, 255, 255))
                # 把原图像粘贴到背景图片上，忽略透明度
                background.paste(img, mask=img.split()[3])  # 3 是 alpha 通道的索引
                img = background
            
            elif img.mode == 'P':
                img = img.convert('RGB')
            # 将调整大小后的图片转换为 JPEG 格式并存入字节流
            with io.BytesIO() as output_buffer:
                img.save(output_buffer, format='JPEG', quality=85)  # 可调整quality参数来优化大小
                jpeg_data = output_buffer.getvalue()

            # 将 JPEG 格式的图片内容作为字节流添加到form_data中
            if isinstance(original_image,str):
                form_data.add_field('original_image', jpeg_data, 
                                filename=original_image+".jpeg", 
                                content_type='image/jpeg')
            else:
                form_data.add_field('original_image', jpeg_data, 
                                    filename=original_image.filename, 
                                    content_type='image/jpeg')
            # # 读取文件内容作为字节流，并传给`FormData`
            # form_data.add_field('original_image', original_image.file.read(), 
            #                     filename=original_image.filename, 
            #                     content_type=original_image.content_type)
    # if mask_image:
        
    #     # 确保光标回到文件的开头
    #     mask_image.file.seek(0)
    #     # 读取文件内容作为字节流，并传给`FormData`
    #     form_data.add_field('mask_image', mask_image.file.read(), 
    #                         filename=mask_image.filename, 
    #                         content_type=mask_image.content_type)
    # 构造请求的URL
    url = "https://zaiwen.superatmas.com:50008/img_generate"
    
    async with aiohttp.ClientSession() as session:
        # 设置20秒超时
        timeout = ClientTimeout(total=20)

        async with session.post(url, data=form_data, timeout=timeout) as response:
            response_data = await response.json()
        print(response_data)
        return response_data
    
@draw_router.get("/img_generate_task/{task_id}")
async def forward_img_generate_task_status(task_id: str):
    global img_url
    if task_id=="hasexceeded":
        img = Image.open('buy.jpg')
        img_io = io.BytesIO()
        img.save(img_io, format='JPEG')
        img_io.seek(0)
        return StreamingResponse(img_io, media_type="image/jpeg")
    # 判断task_id是否含有main-:
    if task_id.startswith("main-"):
        # url是一张图片，返回这张图片
        url = "https://qph.cf2.poecdn.net/"+task_id
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
        img_data = response.content
        
        # 将图片数据转换为WebP格式
        image = Image.open(io.BytesIO(img_data))
        webp_io = io.BytesIO()
        image.save(webp_io, format="JPEG")
        webp_io.seek(0)
        return StreamingResponse(webp_io, media_type="image/jpeg")
    if task_id.startswith("org-"):
        url = img_url[task_id]
        # delete the url
        del img_url[task_id]
        
        response = requests.get(url)
        img_data = response.content
        
        image = Image.open(io.BytesIO(img_data))
        webp_io = io.BytesIO()
        image.save(webp_io, format="JPEG")
        webp_io.seek(0)
        return StreamingResponse(webp_io, media_type="image/jpeg")
    if task_id.startswith("anytext-"):
        print(f"task_id:{task_id}")
        url = img_url[task_id]
        print("url")
        if url:
            image = Image.open(url)
            webp_io = io.BytesIO()
            image.save(webp_io, format="JPEG")
            webp_io.seek(0)
            return StreamingResponse(webp_io, media_type="image/jpeg")

    timeout = Timeout(30.0)
    async with httpx.AsyncClient(timeout = timeout) as client:
        response = await client.get(f"img_server_address/img_generate_task/{task_id}")
    return StreamingResponse(response.aiter_bytes(), media_type=response.headers["Content-Type"], headers=dict(response.headers), status_code=response.status_code)




